package com.huan.kafka.consumer;

import com.huan.kafka.producer.KafkaProducer;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.springframework.kafka.annotation.KafkaListener;
import org.springframework.kafka.support.Acknowledgment;
import org.springframework.stereotype.Component;
import org.springframework.util.function.ThrowingFunction;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.*;

/**
 * 多线程消息消费者
 */
@Component
public class MultiThreadKafkaConsumer {
    private ThreadPoolExecutor executor;

    @PostConstruct
    public void init() {
        executor = new ThreadPoolExecutor(5, 5, 1, TimeUnit.MILLISECONDS, new ArrayBlockingQueue<>(100), new ThreadPoolExecutor.CallerRunsPolicy());
    }

    // 加 @PreDestroy，Spring 容器关闭时自动关闭线程池
    @PreDestroy
    public void destroy() {
        // 优雅关闭：等待 30 秒，未完成的任务继续执行，之后强制关闭
        if (executor != null && !executor.isShutdown()) {
            executor.shutdown();
            try {
                if (!executor.awaitTermination(30, TimeUnit.SECONDS)) {
                    executor.shutdownNow(); // 强制关闭未完成任务
                }
            } catch (InterruptedException e) {
                executor.shutdownNow();
                Thread.currentThread().interrupt(); // 保留中断状态
            }
        }
    }

    @KafkaListener(topics = KafkaProducer.TOPIC, groupId = "kafka-springboot-001")
    public void consumer(List<ConsumerRecord<String, String>> records, Acknowledgment ack) throws InterruptedException {


        System.out.println("线程 " + Thread.currentThread().getName() + " 收到一个批次，共 " + records.size() + " 条消息。");

        List<CompletableFuture<Void>> completableFutures = new ArrayList<>(10);

        // 在这里遍历处理批次中的每一条消息
        for (ConsumerRecord<String, String> record : records) {
            CompletableFuture<Void> voidCompletableFuture = CompletableFuture.runAsync(() -> {
                System.out.println(Thread.currentThread().getName() + " ==> 处理消息: Topic=" + record.topic() + ", Partition=" + record.partition() + ", Offset=" + record.offset() + ", Value=" + record.value());
                // 模拟业务处理
                try {
                    TimeUnit.MILLISECONDS.sleep(10);
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                }
            }, executor).exceptionallyAsync((ThrowingFunction<Throwable, Void>) throwable -> {
                System.err.println("发生异常:" + throwable.getMessage());
                return null;
            });

            completableFutures.add(voidCompletableFuture);

        }

        try {
            CompletableFuture.allOf(completableFutures.toArray(new CompletableFuture[0]))
                    // 阻塞直到所有任务完成
                    .get();
        } catch (ExecutionException e) {
            throw new RuntimeException(e);
        }

        System.out.println("批次处理完成，手动提交 Ack。");
        // 当整个批次都处理完成后，手动提交
        ack.acknowledge();
    }
}
